import { Embeddings, EmbeddingsParams } from "@langchain/core/embeddings";
import axios from "axios";

export interface DouBaoEmbeddingsParams extends EmbeddingsParams {
    key: string;
    url: string;
    model: string;
}

export class DouBaoEmbeddings extends Embeddings {
    private key: string;
    private url: string;
    private model: string;
    constructor(params?: Partial<EmbeddingsParams> & DouBaoEmbeddingsParams) {
        super(params);
        this.key = params?.key || "";
        this.url = params?.url || "";
        this.model = params?.model || "";
    }
    async embedDocuments(documents: string[]): Promise<number[][]> {
        try {
            const res = await axios.post(
                this.url,
                {
                    encoding_format: "float",
                    input: documents,
                    model: this.model
                },
                {
                    headers: {
                        "Content-Type": "application/json",
                        Authorization: `Bearer ${this.key}`
                    }
                }
            );
            // console.log(res);

            let result = [];
            for (let item of res.data.data) {
                result.push(item.embedding);
            }
            return result;
        } catch (error) {
            console.log(error); 
            return [];
        }


    }
    embedQuery(document: string): Promise<number[]> {
        throw new Error("Method not implemented.");
    }
}
